import numpy as np
from .sampler import ClassAwareSampler

import torch
import torchvision
from torchvision import transforms
import torchvision.datasets
from datasets.other import get_imb_data, Imagenet32, noisify_cifar100_asymmetric


class IMBALANCECIFAR10(torchvision.datasets.CIFAR10):
    cls_num = 10

    def __init__(self, root='~/data', train=True, transform=None,
                 r_ood=0.2, r_id=0.2, r_imb=0.1, seed=0, asym=False):
        super().__init__(root, train=train, transform=transform)
        if train is False:
            return
        if r_imb > 0.:
            self.data, self.targets = get_imb_data(self.data, self.targets,
                                                   self.cls_num, r_imb, seed)
            print('built imbalanced dataset')
        np.random.seed(seed)
        if r_ood > 0.:
            ids_ood = [i for i in range(len(self.targets)) if np.random.random() < r_ood]
            imagenet32 = Imagenet32(root='~/data/imagenet32', train=True)
            img_ood = imagenet32.data[np.random.permutation(range(len(imagenet32)))[:len(ids_ood)]]
            self.ids_ood = ids_ood
            self.data[ids_ood] = img_ood
            print('Mixing in OOD noise')

        if r_id > 0.:
            ids_not_ood = [i for i in range(len(self.targets)) if i not in ids_ood]
            ids_id = [i for i in ids_not_ood if np.random.random() < (r_id / (1 - r_ood))]
            if asym:
                if self.cls_num == 10:
                    transition = {0: 0, 2: 0, 4: 7, 7: 7, 1: 1, 9: 1, 3: 5, 5: 3, 6: 6, 8: 8}
                    for i, t in enumerate(self.targets):
                        if i in ids_id:
                            self.targets[i] = transition[t]
                else:
                    self.targets = noisify_cifar100_asymmetric(self.targets, r_id, seed)
                    self.targets = torch.tensor(self.targets, dtype=torch.long)
                print('Mixing in ID asym noise')
            else:
                for i, t in enumerate(self.targets):
                    if i in ids_id:
                        self.targets[i] = int(np.random.random() * self.cls_num)
                print('Mixing in ID noise')
            self.ids_id = ids_id


class CIFAR10_LT(object):

    def __init__(self, distributed, r_id, r_ood, asym,
                 imb_factor=0.01, batch_size=128, num_works=40):
        train_transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

        eval_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])

        train_dataset = IMBALANCECIFAR10(r_id=r_id, r_ood=r_ood, r_imb=imb_factor, asym=asym, seed=0, train=True,
                                         transform=train_transform)
        eval_dataset = torchvision.datasets.CIFAR10(root='~/data', train=False, download=True, transform=eval_transform)

        self.cls_num_list = np.unique(train_dataset.targets, return_counts=True)[1]

        self.dist_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if distributed else None
        self.train_instance = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=batch_size, shuffle=True,
            num_workers=num_works, pin_memory=True, sampler=self.dist_sampler)

        balance_sampler = ClassAwareSampler(train_dataset)
        self.train_balance = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=batch_size, shuffle=False,
            num_workers=num_works, pin_memory=True, sampler=balance_sampler)

        self.eval = torch.utils.data.DataLoader(
            eval_dataset,
            batch_size=batch_size, shuffle=False,
            num_workers=num_works, pin_memory=True)
